#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
import os
import librosa
import torch
from model import SenseVoiceSmall
from funasr.utils.postprocess_utils import rich_transcription_postprocess

# 1. 初始化模型
model_dir = "model/iic/SenseVoiceSmall"
m, kwargs = SenseVoiceSmall.from_pretrained(model=model_dir, device="设备") # 自行指定设备
m.eval()

# 2. 音频加载函数（转换为模型可处理的张量）
def load_audio(file_path, sr=16000):
    audio, _ = librosa.load(file_path, sr=sr)
    audio_tensor = torch.from_numpy(audio).unsqueeze(0)  # 增加批次维度 [1, T]
    return audio_tensor.to(device=kwargs["device"])  # 转移到模型设备

# 3. 配置批量处理参数（根据你的实际路径修改）
audio_dir = "音频目录"  # 音频文件所在目录
result_dir = "结果保存目录"  # 结果保存目录
output_file = os.path.join(result_dir, "结果txt名")  # 完整结果文件路径
start_idx = 起始  # 起始索引（test1.mp3）
end_idx = 结束  # 结束索引（test100.mp3）
file_ext = ".mp3"  # 音频文件后缀

# 4. 确保结果目录存在（若不存在则创建）
os.makedirs(result_dir, exist_ok=True)

# 5. 批量处理并写入结果
with open(output_file, "w", encoding="utf-8") as f:
    for i in range(start_idx, end_idx + 1):
        # 生成文件名（如test1.mp3、test2.mp3...）
        audio_name = f"test{i}{file_ext}"
        audio_path = os.path.join(audio_dir, audio_name)
        
        # 检查文件是否存在
        if not os.path.exists(audio_path):
            print(f"警告：文件不存在 - {audio_path}")
            f.write(f"{i}. [文件不存在]\n")
            continue
        
        # 加载音频并推理
        try:
            audio_tensor = load_audio(audio_path)
            res = m.inference(
                data_in=audio_tensor,
                data_type="tensor",
                language="auto",
                use_itn=False,
                ban_emo_unk=False,
                **kwargs,
            )
            # 提取并处理文本结果
            text = rich_transcription_postprocess(res[0][0]["text"])
            # 按索引格式写入文件（如"1. 文本内容"）
            f.write(f"{i}. {text}\n")
            print(f"已处理：{audio_name} -> {text}")  # 控制台打印进度
        except Exception as e:
            print(f"处理失败 {audio_name}：{str(e)}")
            f.write(f"{i}. [处理失败：{str(e)}]\n")  # 失败时记录错误信息

print(f"批量处理完成，结果已保存至：{output_file}")